#include "KeyGenerator.h"
#include "SymmetricKeyMemory.h"

#include <QCryptographicHash>
#include <QDate>

#include <botan/botan.h>
#include <botan/rsa.h>
#include <botan/x509cert.h>
#include <botan/x509self.h>
#include <botan/pem.h>
#include <botan/exceptn.h>
#include <botan/internal/es_win32.h>
#include <botan/oids.h>

#include <assert.h>

using namespace kryptos::entities;

namespace kryptos { namespace cryptography {

bool initialized = false;

KeyGenerator::KeyGenerator()
{
}

void KeyGenerator::init()
{
    if (!initialized)
    {
        Botan::LibraryInitializer init;
        initialized = true;
    }
}

PrivateKey KeyGenerator::encodeKey(const PrivateKey& pk)
{
    assert(initialized);

    Botan::AutoSeeded_RNG rng;
    Botan::AlgorithmIdentifier algoId(Botan::OIDS::lookup("RSA"), Botan::AlgorithmIdentifier::USE_NULL_PARAM);
    Botan::MemoryVector<Botan::byte> privBytes((const unsigned char*)pk.key().constData(), pk.key().size());
    Botan::MemoryVector<Botan::byte> pubBytes((const unsigned char*)pk.publicKeyPart().key().constData(), pk.publicKeyPart().key().size());

    Botan::RSA_PrivateKey privkey(algoId, privBytes, rng);
    Botan::RSA_PublicKey  pubkey(algoId, pubBytes);

    Botan::SecureVector<Botan::byte> publicBytes  = Botan::X509::BER_encode(pubkey);
    Botan::SecureVector<Botan::byte> privateBytes = Botan::PKCS8::BER_encode(privkey, rng, pk.info().passphrase());

    QByteArray publicArray((char*)publicBytes.begin(), publicBytes.size());
    QByteArray privateArray((char*)privateBytes.begin(), privateBytes.size());

    return PrivateKey(privateArray, publicArray, pk.info());
}

PrivateKey KeyGenerator::generateKey(const KeyInfo& info, KeyInfo::KeyType type)
{
    assert(initialized);

    Botan::AutoSeeded_RNG rng;
    Botan::Win32_EntropySource entropySource;

    rng.add_entropy_source(&entropySource);
    rng.reseed(10096);

    Botan::RSA_PrivateKey rsaPrivate(rng, info.size());

    QByteArray privateArray((char*)rsaPrivate.pkcs8_private_key().begin(), rsaPrivate.pkcs8_private_key().size());
    QByteArray publicArray((char*)rsaPrivate.x509_subject_public_key().begin(), rsaPrivate.x509_subject_public_key().size());

    return PrivateKey(privateArray, publicArray,info);
}

SymmetricKey* KeyGenerator::generateSessionKey(int size)
{
    Botan::AutoSeeded_RNG rng;
    Botan::Win32_EntropySource entropySource;

    rng.add_entropy_source(&entropySource);
    rng.reseed(10096);

    SymmetricKeyMemory* sessionKey = new SymmetricKeyMemory((const char*)rng.random_vec(size).begin(), size, true);

    return sessionKey;
}

PrivateKey KeyGenerator::generateEncodedKey(const KeyInfo& info, KeyInfo::KeyType type)
{
    assert(initialized);

    Botan::AutoSeeded_RNG rng;
    Botan::Win32_EntropySource entropySource;

    rng.add_entropy_source(&entropySource);
    rng.reseed(10096);

    Botan::RSA_PrivateKey rsaPrivate(rng, info.size());

    Botan::SecureVector<Botan::byte> publicBytes  = Botan::X509::BER_encode(rsaPrivate);
    Botan::SecureVector<Botan::byte> privateBytes = Botan::PKCS8::BER_encode(rsaPrivate, rng, info.passphrase());

    QByteArray publicArray((char*)publicBytes.begin(), publicBytes.size());
    QByteArray privateArray((char*)privateBytes.begin(), privateBytes.size());

    return PrivateKey(privateArray, publicArray,info);
}


PrivateKey KeyGenerator::privateKeyFromBytes(const KeyInfo& info, const std::vector<char>& privatePart, const std::vector<char>& publicPart)
{
    assert(initialized);

    Botan::AutoSeeded_RNG rng;
    Botan::DataSource_Memory privateSource((const Botan::byte*)&privatePart[0], privatePart.size());
    Botan::Private_Key* rsaprivate = Botan::PKCS8::load_key(privateSource, rng, info.passphrase());

    Botan::DataSource_Memory publicSource((const Botan::byte*)&publicPart[0], publicPart.size());
    Botan::Public_Key* rsapublic = Botan::X509::load_key(publicSource);

    QByteArray privateArray((char*)rsaprivate->pkcs8_private_key().begin(), rsaprivate->pkcs8_private_key().size());
    QByteArray publicArray((char*)rsapublic->x509_subject_public_key().begin(), rsapublic->x509_subject_public_key().size());

    return PrivateKey(privateArray, publicArray, info);
}

PublicKey KeyGenerator::publicKeyFromBytes(const KeyInfo& info, const std::vector<char>& publicPart)
{
    assert(initialized);

    Botan::AutoSeeded_RNG rng;
    Botan::DataSource_Memory publicSource((const Botan::byte*)&publicPart[0], publicPart.size());
    Botan::Public_Key* rsapublic = Botan::X509::load_key(publicSource);

    QByteArray publicArray((char*)rsapublic->x509_subject_public_key().begin(), rsapublic->x509_subject_public_key().size());

    return PublicKey(publicArray, info);
}

std::string KeyGenerator::generateKeyId(const PrivateKey& pk)
{
    assert(initialized);

    QCryptographicHash shaHash(QCryptographicHash::Sha1);
    shaHash.addData(pk.info().name().c_str());
    shaHash.addData(pk.info().email().c_str());
    shaHash.addData(pk.key());

    return shaHash.result().toHex().data();
}

std::string KeyGenerator::certificateFromKey(const PrivateKey& pk)
{
    assert(initialized);

    Botan::AutoSeeded_RNG rng;
    Botan::AlgorithmIdentifier algoId(Botan::OIDS::lookup("RSA"), Botan::AlgorithmIdentifier::USE_NULL_PARAM);
    Botan::MemoryVector<Botan::byte> bytes((const unsigned char*)pk.key().constData(), pk.key().size());

    Botan::RSA_PrivateKey privkey(algoId, bytes, rng);

    Botan::X509_Cert_Options opt;
    opt.common_name = pk.info().name();
    opt.country     = "54";
    opt.email       = pk.info().email();

    Botan::X509_Certificate cert = Botan::X509::create_self_signed_cert(opt, privkey, "SHA-256", rng);
    return cert.PEM_encode();
}

PublicKey KeyGenerator::publicKeyFromCertificate(const std::string& certificateFile)
{
    assert(initialized);

    Botan::X509_Certificate cert1(certificateFile);
    Botan::Public_Key* publicKey = cert1.subject_public_key();
    std::string name  = cert1.subject_info("Name").at(0);
    std::string email = cert1.subject_info("Email").at(0);

    KeyInfo keyinfo(KeyInfo::RSA, name, email, "imported", " ", 0);

    QByteArray publicArray((char*)publicKey->x509_subject_public_key().begin(), publicKey->x509_subject_public_key().size());

    keyinfo.setKeyid(QCryptographicHash::hash(publicArray, QCryptographicHash::Sha1).toHex().data());
    keyinfo.setPassphrase(QCryptographicHash::hash(QByteArray(keyinfo.passphrase().c_str()), QCryptographicHash::Sha1).data());
    keyinfo.setTimestamp(QDate::currentDate().toString().toStdString());

    return PublicKey(publicArray, keyinfo);
}
}}

